# ===----------------------------------------------------------------------=== #
# Copyright (c) 2025, Modular Inc. All rights reserved.
#
# Licensed under the Apache License v2.0 with LLVM Exceptions:
# https://llvm.org/LICENSE.txt
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ===----------------------------------------------------------------------=== #
"""Audio generation interface definitions for Modular's MAX API.

This module provides data structures and interfaces for handling audio generation
responses, including status tracking and audio data encapsulation.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Generic

import msgspec
import numpy as np
import numpy.typing as npt
from max.interfaces.context import BaseContext, SamplingParams
from max.interfaces.pipeline import PipelineInputs, PipelineOutput
from max.interfaces.request import Request, RequestID
from max.interfaces.status import GenerationStatus
from typing_extensions import TypeVar


@dataclass(frozen=True)
class AudioGenerationRequest(Request):
    model: str = field()
    """The name of the model to be used for generating audio chunks. This should match
    the available models on the server and determines the behavior and
    capabilities of the response generation.
    """

    input: str | None = None
    """The text to generate audio for. The maximum length is 4096 characters.
    """

    audio_prompt_tokens: list[int] = field(default_factory=list)
    """The prompt speech IDs to use for audio generation."""

    audio_prompt_transcription: str = ""
    """The audio prompt transcription to use for audio generation."""

    sampling_params: SamplingParams = field(default_factory=SamplingParams)
    """Request sampling configuration options."""

    _assistant_message_override: str | None = None
    """(ONLY FOR BENCHMARKING PURPOSES) An assistant message that replaces the
    speech token pattern."""

    prompt: list[int] | str | None = field(default=None)
    """Optionally provide a preprocessed list of token ids or a prompt string to pass as input directly into the model.
    This replaces automatically generating TokenGeneratorRequestMessages given the input, audio prompt tokens,
    audio prompt transcription fields."""

    streaming: bool = True
    """Whether to stream the audio generation."""

    buffer_speech_tokens: npt.NDArray[np.integer[Any]] | None = None
    """An optional field potentially containing the last N speech tokens
    generated by the model from a previous request.

    When this field is specified, this tensor is used to buffer the tokens sent
    to the audio decoder.
    """

    def __post_init__(self) -> None:
        if self.prompt is None and self.input is None:
            raise RuntimeError("either token_ids or input must be provided.")


class AudioGenerationMetadata(
    msgspec.Struct, tag=True, omit_defaults=True, kw_only=True
):
    """
    Represents metadata associated with audio generation.

    This class will eventually replace the metadata dictionary used throughout
    the AudioGenerationOutput object, providing a structured and type-safe
    alternative for audio generation metadata.

    Configuration:
        sample_rate: The sample rate of the generated audio in Hz.
        duration: The duration of the generated audio in seconds.
        chunk_id: Identifier for the audio chunk (useful for streaming).
        timestamp: Timestamp when the audio was generated.
        final_chunk: Whether this is the final chunk in a streaming sequence.
        model_name: Name of the model used for generation.
        request_id: Unique identifier for the generation request.
        tokens_generated: Number of tokens generated for this audio.
        processing_time: Time taken to process this audio chunk in seconds.
        echo: Echo of the input prompt or identifier for verification.
    """

    sample_rate: int | None = None
    duration: float | None = None
    chunk_id: int | None = None
    timestamp: str | None = None
    final_chunk: bool | None = None
    model_name: str | None = None
    request_id: RequestID | None = None
    tokens_generated: int | None = None
    processing_time: float | None = None
    echo: str | None = None

    def to_dict(self) -> dict[str, int | float | str | bool]:
        """
        Convert the metadata to a dictionary format.

        Returns:
            dict[str, any]: Dictionary representation of the metadata.
        """
        result = {}
        for attr in self.__annotations__:
            if value := getattr(self, attr, None):
                result[attr] = value
        return result


AudioGenerationContextType = TypeVar(
    "AudioGenerationContextType", bound=BaseContext
)
"""Type variable for audio generation context types.

This type variable is bound to BaseContext and represents the specific context
type used in audio generation pipelines. It allows for type-safe generic
programming while ensuring that all context types inherit from BaseContext
and maintain the required interface for audio generation operations.
"""


@dataclass(frozen=True)
class AudioGenerationInputs(
    PipelineInputs, Generic[AudioGenerationContextType]
):
    """Input data structure for audio generation pipelines.

    This class represents the input data required for audio generation operations
    within the pipeline framework. It extends PipelineInputs and provides type-safe
    generic support for different audio generation context types.
    """

    batch: dict[RequestID, AudioGenerationContextType]
    """A dictionary mapping RequestID to AudioGenerationContextType instances.
    This batch structure allows for processing multiple audio generation
    requests simultaneously while maintaining request-specific context
    and configuration data.
    """


class AudioGenerationOutput(msgspec.Struct, tag=True, omit_defaults=True):
    """Represents a response from the audio generation API.

    This class encapsulates the result of an audio generation request, including
    the final status, generated audio data, and optional buffered speech tokens.
    """

    final_status: GenerationStatus
    """The final status of the generation process."""
    steps_executed: int
    """The number of steps previously executed."""
    audio_data: npt.NDArray[np.float32] = msgspec.field(
        default_factory=lambda: np.array([], dtype=np.float32)
    )
    """The generated audio data, if available."""
    buffer_speech_tokens: npt.NDArray[np.integer[Any]] | None = None
    """Buffered speech tokens, if available."""
    metadata: AudioGenerationMetadata = msgspec.field(
        default_factory=AudioGenerationMetadata
    )
    """Metadata associated with the audio generation, such as chunk information, prompt details, or other relevant context."""

    @property
    def is_done(self) -> bool:
        """Indicates whether the audio generation process is complete.

        Returns:
            :class:`bool`: ``True`` if generation is done, ``False`` otherwise.
        """
        return self.final_status.is_done


def _check_audio_generator_output_implements_pipeline_output(
    x: AudioGenerationOutput,
) -> PipelineOutput:
    return x
